#!/usr/bin/env python
# -*- coding: utf-8 -*-
#

import os
import sys

src_x64 = bytearray.fromhex("4803C3488B0848394108" "7522" "488B5008483902" "7519" "48890a")
dst_x64 = bytearray.fromhex("4803C3488B0848394108" "9090" "488B5008483902" "9090" "48890a")

# WDK 10.0.19041.0
src_arm64 = bytearray.fromhex("880540F91F0109EB" "21020054" "2A0540F9480140F91F0109EB" "A1010054" "4C0100F98A0500F9")
dst_arm64 = bytearray.fromhex("880540F91F0109EB" "1F2003D5" "2A0540F9480140F91F0109EB" "1F2003D5" "4C0100F98A0500F9")

# WDK 10.0.22621.0
src_22621_arm64 = bytearray.fromhex("880540F91F0109EB" "C1010054" "2A0540F9480140F91F0109EB" "41010054" "4C0100F960120091")
dst_22621_arm64 = bytearray.fromhex("880540F91F0109EB" "1F2003D5" "2A0540F9480140F91F0109EB" "1F2003D5" "4C0100F960120091")

if len(sys.argv) < 3:
    raise RuntimeError("usage: patch x64|aarch64 path_to_hwfs.sys|wdk_ver")

hwfs_sys = sys.argv[2]

with open(hwfs_sys, 'rb') as file:
    bin_data = file.read()

if sys.argv[1] == 'x64':
    src = src_x64
    dst = dst_x64
elif sys.argv[1] == 'ARM64':
    src = src_arm64
    dst = dst_arm64
    if len(sys.argv) >= 4:
        if sys.argv[3] == "10.0.22621.0":
            print("using wdk 2261")
            src = src_22621_arm64
            dst = dst_22621_arm64
    else:
        # MSBuild EWDK 22000 doesn't fill arg about WDK ver. Detect manually.
        if bin_data.count(src_22621_arm64) or bin_data.count(dst_22621_arm64):
            src = src_22621_arm64
            dst = dst_22621_arm64
else:
    raise RuntimeError("Not supported architecture {sys.argv[1]}")

count = bin_data.count(src)
if count == 0:
    dst_count = bin_data.count(dst)
    if dst_count == 1:
        print("Already patched.")
        sys.exit(0)
    raise RuntimeError("Pattern not found")

if count != 1:
    raise RuntimeError("Pattern found more then once")

with os.fdopen(os.open(hwfs_sys, os.O_WRONLY | os.O_CREAT, 0o666), "wb") as file:
    file.write(bin_data.replace(src, dst))
